 package com.gmrz.uaf.crypto;

 import com.gmrz.uaf.common.UAFRuntimeException;
 import com.gmrz.uaf.crypto.internal.keyFormats.UAFAlgKeyECCX962Raw;
 import com.gmrz.uaf.crypto.spi.CryptoService;
import com.gmrz.uaf.crypto.spi.UAFAuthAlgorithmSuite;
import com.gmrz.uaf.crypto.spi.UAFKeyFormat;
import com.gmrz.uaf.crypto.spi.handlers.*;
import com.gmrz.uaf.protocol.v1.processor.exception.UAFErrorCode;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.Provider;
import java.security.PublicKey;
import java.security.Security;
import java.security.cert.Certificate;
import java.util.*;

 /**
  * 加密引擎
  */
 @Singleton
 public class CryptoEngine
 {
   private static final Logger LOG = LogManager.getLogger(CryptoEngine.class);
 
   private Map<Integer, UAFAuthAlgorithmSuite> authAlgos = new HashMap<Integer, UAFAuthAlgorithmSuite>();
   private Map<Integer, UAFKeyFormat> keyFormats = new HashMap<Integer, UAFKeyFormat>();
 
   private Map<String, CipherHandler> cipherHandlers = new HashMap<String, CipherHandler>();
   private Map<String, HMACHandler> hmacHandlers = new HashMap<String, HMACHandler>();
   private Map<String, HashHandler> hashHandlers = new HashMap<String, HashHandler>();
   private AttestationCertHandler attestationCertHandler = null;
   private CryptoConfig cryptoConfig = null;
 
   @Inject
   public CryptoEngine(CryptoConfig cryptoConfig) {
     LOG.info("Creating instance of CryptoEngine");
     if (cryptoConfig == null) {
       String errorMessage = "CryptoConfig instance cannot be null.";
       LOG.error(errorMessage);
       throw new UAFRuntimeException(errorMessage);
     }
     this.cryptoConfig = cryptoConfig;
 
     addProviders();
 
     loadCustomCryptoService(cryptoConfig.getCustomProperties());
   }
 
   public synchronized void initialize(List<String> allowedOperations)
   {
     LOG.debug("Allowed Operations to be used in the initialization - {}", Arrays.toString(allowedOperations.toArray()));
 
     List<String> defaultOperations = getDefaultOperations(allowedOperations);
     loadDefaultCryptoService(defaultOperations, this.cryptoConfig);
   }
 
   private void loadDefaultCryptoService(List<String> defaultOperations, CryptoConfig cryptoConfig)
   {
     CryptoService defaultCryptoService = new DefaultCryptoService(defaultOperations, cryptoConfig);
     loadCryptoService(defaultCryptoService);
   }
 
   private void loadCryptoService(CryptoService cryptoService) {
     LOG.info("BEGIN - Loading CryptoService class - {}", cryptoService.getClass().toString());
     initializeHandlers(cryptoService);
     initializeSignatureAlgos(cryptoService);
     initializeKeyFormats(cryptoService);
     LOG.info("END - Loading CryptoService class - {}", cryptoService.getClass().toString());
   }
 
   private List<String> getDefaultOperations(List<String> allowedOperations)
   {
     List<String> operations = new ArrayList<String>();
 
     if ((!this.cipherHandlers.containsKey(Operations.SERVER_DATA_CIPHER)) && (allowedOperations.contains(Operations.SERVER_DATA_CIPHER))) {
       LOG.info("There is no custom CipherHandler implementation for operation - {}", Operations.SERVER_DATA_CIPHER);
       operations.add(Operations.SERVER_DATA_CIPHER);
     }
     if ((!this.cipherHandlers.containsKey(Operations.PASSWORD_CIPHER)) && (allowedOperations.contains(Operations.PASSWORD_CIPHER))) {
       LOG.info("There is no custom CipherHandler implementation for operation - {}", Operations.PASSWORD_CIPHER);
       operations.add(Operations.PASSWORD_CIPHER);
     }
     if ((!this.hashHandlers.containsKey(Operations.DATA_HASH)) && (allowedOperations.contains(Operations.DATA_HASH))) {
       LOG.info("There is no custom data hash implementation for operation - {}", Operations.DATA_HASH);
       operations.add(Operations.DATA_HASH);
     }
     if ((!this.hmacHandlers.containsKey(Operations.HMAC)) && (allowedOperations.contains(Operations.HMAC))) {
       LOG.info("There is no custom HMAC implementation for operation - {}", Operations.HMAC);
       operations.add(Operations.HMAC);
     }
     return operations;
   }
 
   private <String, T> void insertHandler(Map<String, T> handlers, Map<String, T> internalHandlers, String logPrefix)
   {
     if (handlers == null) {
       return;
     }
     for (Iterator<String> i$ = handlers.keySet().iterator(); i$.hasNext(); ) { 
    	 String key = i$.next();
 
       if (!internalHandlers.containsKey(key)) {
         LOG.info("{} implementation with operation - {}", logPrefix, key);
         internalHandlers.put(key, handlers.get(key));
       } }
   }
 
   public UAFKeyFormat getKeyFormat(int keyFormat)
     throws KeyFormatNotFoundException
   {
     UAFKeyFormat uafKeyFormat = this.keyFormats.get(Integer.valueOf(keyFormat));
     LOG.info("{} uafKeyFormat getName",uafKeyFormat.getName());
     if (uafKeyFormat == null) {
       String errorMsg = "Unsupported key format - " + keyFormat;
       LOG.error(errorMsg);
       throw new KeyFormatNotFoundException(errorMsg);
     }
     return uafKeyFormat;
   }
 
   private void loadCustomCryptoService(Properties uafProperties) {
     String cryptoCustomClassName = this.cryptoConfig.getCryptoServiceClassName();
     if ((cryptoCustomClassName != null) && (!cryptoCustomClassName.trim().isEmpty()))
       try {
         CryptoService customCryptoService = (CryptoService)Class.forName(cryptoCustomClassName).getConstructor(new Class[0]).newInstance();
 
         LOG.info("Using custom crypto plugin implementation : " + customCryptoService.getClass().getName());
 
         customCryptoService.initialize(uafProperties);
 
         loadCryptoService(customCryptoService);
       }
       catch (Exception e) {
         String errorMessage = "Failed to load the custom crypto service class";
         LOG.error("{} - {} due to exception - {}", errorMessage, cryptoCustomClassName, e.getMessage());
         throw new UAFRuntimeException(errorMessage, e);
       }
   }
 
   private void initializeKeyFormats(CryptoService cryptoService)
   {
     List<UAFKeyFormat> keyFormatsList = cryptoService.getSupportedKeyFormats();
     if ((keyFormatsList == null) || (keyFormatsList.isEmpty())) {
       return;
     }
 
     for (UAFKeyFormat keyFormat : keyFormatsList) {
       int format = keyFormat.getFormat();
 
       if (!this.keyFormats.containsKey(Integer.valueOf(format))) {
         LOG.info("Loading UAFKeyFormat with key format - {} and name - {}", Integer.valueOf(format), keyFormat.getName());
 
         this.keyFormats.put(Integer.valueOf(format), keyFormat);
       }
     }
   }
 
   private void initializeSignatureAlgos(CryptoService cryptoService)
   {
     List<UAFAuthAlgorithmSuite> signatureAlgos = cryptoService.getSupportedAlgos();
     if ((signatureAlgos == null) || (signatureAlgos.isEmpty())) {
       return;
     }
 
     for (UAFAuthAlgorithmSuite uafAuthAlgorithmSuite : signatureAlgos)
       insertSignatureAlgorithm(uafAuthAlgorithmSuite);
   }
 
   private void insertSignatureAlgorithm(UAFAuthAlgorithmSuite uafAuthAlgorithmSuite)
   {
     if (!this.authAlgos.containsKey(Integer.valueOf(uafAuthAlgorithmSuite.getSuite()))) {
       LOG.info("Loading UAFAuthAlgorithmSuite with algorithm suite - {} and name - {}", Integer.valueOf(uafAuthAlgorithmSuite.getSuite()), uafAuthAlgorithmSuite.getName());
       if (uafAuthAlgorithmSuite.init()) {
         this.authAlgos.put(Integer.valueOf(uafAuthAlgorithmSuite.getSuite()), uafAuthAlgorithmSuite);
       }
       else
       {
         LOG.error("Initialization failed for UAF Auth Algorithm - {} with name - {}.Hence this algorithm will be unavailable on the server.", Integer.valueOf(uafAuthAlgorithmSuite.getSuite()), uafAuthAlgorithmSuite.getName());
       }
     }
   }
 
   private void initializeHandlers(CryptoService cryptoService)
   {
     LOG.info("Loading CipherHandler implementations -  ");
     insertHandler(cryptoService.getCipherHandlers(), this.cipherHandlers, "Loading CipherHandler ");
     LOG.info("Loading HashHandler implementations -  ");
     insertHandler(cryptoService.getHashHandlers(), this.hashHandlers, "Loading HashHandler ");
     LOG.info("Loading HMACHandler implementations -  ");
     insertHandler(cryptoService.getHMACHandlers(), this.hmacHandlers, "Loading HMACHandler ");
 
     if (this.attestationCertHandler == null) {
       LOG.info("AttestationCertHandler implementation provided by {} will be used.", cryptoService.getClass().toString());
       this.attestationCertHandler = cryptoService.getAttestationCertHandler();
     }
   }
 
   public boolean verifySignature(int keyFormat, byte[] pubKey, int signAlgo, byte[] signedData, byte[] signature)
     throws AlgorithmNotFoundException, IOException, KeyFormatNotFoundException, GeneralSecurityException
   {
     LOG.debug("Veriying signature using keyFormat - {} and algoSuite - {}", Integer.valueOf(keyFormat), Integer.valueOf(signAlgo));
     UAFAuthAlgorithmSuite authAlgorithmSuite = getAuthAlgorithmSuite(signAlgo);
     LOG.debug("PublicKey begin");
     PublicKey key=null;
     if(signAlgo==1)
     {
       UAFAlgKeyECCX962Raw ECCXkeySpec= new UAFAlgKeyECCX962Raw();
       key=ECCXkeySpec.getKey(pubKey,signAlgo);
     }
     else
     {
       key = getKeyFormat(keyFormat).getKey(pubKey, signAlgo);
     }
     LOG.debug("PublicKey end verifySignature start");
     return authAlgorithmSuite.verifySignature(key, signedData, signature);
   }
 
   public UAFAuthAlgorithmSuite getAuthAlgorithmSuite(int signAlgo)
     throws AlgorithmNotFoundException
   {
     UAFAuthAlgorithmSuite authAlgorithmSuite = this.authAlgos.get(Integer.valueOf(signAlgo));
     if (authAlgorithmSuite == null) {
       String errorMsg = "Unsupported signing algorithm - " + signAlgo;
       LOG.error(errorMsg);
       throw new AlgorithmNotFoundException(UAFErrorCode.PROTOCOL_ALGORITHM_NOT_FOUND, errorMsg);
     }
     return authAlgorithmSuite;
   }
 
   public CipherHandler getServerDataCipherHandler()
   {
     CipherHandler cipherHandler = this.cipherHandlers.get(Operations.SERVER_DATA_CIPHER);
     if (cipherHandler == null) {
       throw new UAFRuntimeException("CipherHandler implementation is not found for server data.");
     }
     return cipherHandler;
   }
 
   public CipherHandler getDBPasswordCipherHandler() {
     CipherHandler cipherHandler = this.cipherHandlers.get(Operations.PASSWORD_CIPHER);
     if (cipherHandler == null) {
       throw new UAFRuntimeException("CipherHandler implementation is not found for database password.");
     }
     return cipherHandler;
   }
 
   public HMACHandler getHmacHandler() {
     HMACHandler hmacHandler = this.hmacHandlers.get(Operations.HMAC);
     if (hmacHandler == null) {
       throw new UAFRuntimeException("CipherHandler implementation is not found for generating HMAC.");
     }
     return hmacHandler;
   }
 
   public HashHandler getDataHashHandler() {
     HashHandler hashHandler = this.hashHandlers.get(Operations.DATA_HASH);
     if (hashHandler == null) {
       throw new UAFRuntimeException("HashHandler implementation is not found for data hashing.");
     }
     return hashHandler;
   }
 
   public boolean verifySignature(int signAlgo, Certificate attestationCert, byte[] signedData, byte[] signature) throws Exception
   {
     // 获取认证算法套件对象（就是签名算法）
     UAFAuthAlgorithmSuite authAlgorithmSuite = getAuthAlgorithmSuite(signAlgo);
     return authAlgorithmSuite.verifySignature(attestationCert.getPublicKey(), signedData, signature);
   }
 
   public AttestationCertHandler getAttestationCertHandler() {
     return this.attestationCertHandler;
   }
 
   public String getAlgorithmName(int signAlgo) throws AlgorithmNotFoundException {
     return getAuthAlgorithmSuite(signAlgo).getName();
   }
 
   private void addProviders()
   {
     String jceProviders = this.cryptoConfig.getJceProviders();
 
     if ((jceProviders == null) || (jceProviders.trim().isEmpty()))
       return;
     try
     {
       String[] providerNames = jceProviders.split(",");
       if ((providerNames != null) && (providerNames.length != 0))
         for (String providerName : providerNames)
           Security.addProvider((Provider)Class.forName(providerName).newInstance());
     }
     catch (ClassNotFoundException e)
     {
       logAndThrowException("Provider Class not found - ", e);
     } catch (InstantiationException e) {
       logAndThrowException("Error while instantiating Providers class - ", e);
     } catch (IllegalAccessException e) {
       logAndThrowException("Invalid Access for Providers class - ", e);
     }
     LOG.info("Providers added successfully.");
   }
 
   private void logAndThrowException(String errorMessage, Exception e) throws UAFRuntimeException
   {
     LOG.error("{} {}", errorMessage, e.getMessage());
     throw new UAFRuntimeException(errorMessage, e);
   }
 }
